from PIL import Image
from PIL import ImageEnhance
import os
import cv2
import numpy as np


# 第一步：进行数据增强，如果数据集不小还是别用了
# 笔记本1060显卡真跑不动，放弃了
# 本代码共采用了四种数据增强，如采用其他数据增强方式，可以参考本代码，随意替换。
# imageDir 为原数据集的存放位置
# saveDir  为数据增强后数据的存放位置

# 翻转图像
def flip(root_path, img_name):
    img = Image.open(os.path.join(root_path, img_name))
    filp_img = img.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
    # filp_img = img.transpose(Image.FLIP_LEFT_RIGHT)
    # filp_img.save(os.path.join(root_path,img_name.split('.')[0] + '_flip.jpg'))
    return filp_img


# 旋转角度
def rotation(root_path, img_name):
    img = Image.open(os.path.join(root_path, img_name))
    rotation_img = img.rotate(20)
    # rotation_img.save(os.path.join(root_path,img_name.split('.')[0] + '_rotation.jpg'))
    return rotation_img


# 随机颜色
def randomColor(root_path, img_name):
    """
    对图像进行颜色抖动
    :param image: PIL的图像image
    :return: 有颜色色差的图像image
    """
    image = Image.open(os.path.join(root_path, img_name))
    random_factor = np.random.randint(0, 31) / 10.  # 随机因子
    color_image = ImageEnhance.Color(image).enhance(random_factor)  # 调整图像的饱和度
    random_factor = np.random.randint(10, 21) / 10.  # 随机因子
    brightness_image = ImageEnhance.Brightness(color_image).enhance(random_factor)  # 调整图像的亮度
    random_factor = np.random.randint(10, 21) / 10.  # 随机因子
    contrast_image = ImageEnhance.Contrast(brightness_image).enhance(random_factor)  # 调整图像对比度
    random_factor = np.random.randint(0, 31) / 10.  # 随机因子
    return ImageEnhance.Sharpness(contrast_image).enhance(random_factor)  # 调整图像锐度


# 对比度增强
def contrastEnhancement(root_path, img_name):
    image = Image.open(os.path.join(root_path, img_name))
    enh_con = ImageEnhance.Contrast(image)
    contrast = 1.5
    image_contrasted = enh_con.enhance(contrast)
    return image_contrasted


# 亮度增强
def brightnessEnhancement(root_path, img_name):
    image = Image.open(os.path.join(root_path, img_name))
    enh_bri = ImageEnhance.Brightness(image)
    brightness = 1.5
    image_brightened = enh_bri.enhance(brightness)
    return image_brightened


# 颜色增强
def colorEnhancement(root_path, img_name):
    image = Image.open(os.path.join(root_path, img_name))
    enh_col = ImageEnhance.Color(image)
    color = 1.5
    image_colored = enh_col.enhance(color)
    return image_colored


# 这里要指定到有图片的目录下,不能指定到数据集目录，这样只能一个一个去执行数据增强了
imageDir = "D://培正文件//机器学习//resnet152识别虫害//resnet152_plant//数据增强" \
           "//Tomato___Tomato_Yellow_Leaf_Curl_Virus"  # 要改变的图片的路径文件夹
saveDir = "D://培正文件//机器学习//resnet152识别虫害//resnet152_plant//增强后的数据" \
          "//Tomato___Tomato_Yellow_Leaf_Curl_Virus_strong"  # 要保存的图片的路径文件夹

for name in os.listdir(imageDir):
    saveName = name[:-4] + "id.jpg"
    image = Image.open(os.path.join(imageDir, name))
    image.save(os.path.join(saveDir, saveName))

    # 亮度增强
    saveName = name[:-4] + "be.jpg"
    saveImage = brightnessEnhancement(imageDir, name)
    saveImage.save(os.path.join(saveDir, saveName))

    # 旋转图像
    saveName = name[:-4] + "fl.jpg"
    saveImage = flip(imageDir, name)
    saveImage.save(os.path.join(saveDir, saveName))

    # 旋转角度
    saveName = name[:-4] + "ro.jpg"
    saveImage = rotation(imageDir, name)
    saveImage.save(os.path.join(saveDir, saveName))
